% This code goes through various levels of disaggregation, based on the
% actual inputs for price stickiness, sector size and I/O linkages
% by Raphael Schoenle, August 10, 2016 and April 2019

% only focus on our fully heterogeneous case

clearvars -except wd
cd(fullfile(wd, 'Figure4_RHS')); 

%number of sectors
n_sectors_full = 341;

%number of periods in IRF
n_periods = 30;
n_periods_idio = 30;

%number of cases
num_case = 18;

%pre-definition of variables
% third dimension: number of shocks: monetary, technology, idiosyncratic
c_imp = zeros(n_periods, num_case, 3);
inf_imp = zeros(n_periods, num_case, 3);
mc_imp = zeros(n_periods, num_case, 3, n_sectors_full);
pos_akshock = zeros(n_sectors_full,1);

%Parameters values
params.rho_a = 0.9;         % persistence agg tech shocks
params.rho_mu = 0.9;         % persistence monetary shock
params.rho_k = 0.9;         % persistence idiosync tech shock
params.beta = 0.9975;      % discount factor
params.theta = 6;          % elast of subst within sectors

% parameter sets: only base case for Figure 1!
delta_set = [0.5];
eta_set = [2];
phi_set = [2 ];
sigma_set = [1];
phi_pi_set = [1.24];
phi_c_set = [0.33]./12;
phi_rho_i_set = [0 ];
phi_gc_i_set = [0];
n_parameters = length(phi_rho_i_set);

% Position indices for shocks
pos_mushock = 1;
pos_ashock = 2;
for i = 1:n_sectors_full
    pos_akshock(i) = 2 + i;
end

% produce aggregation up to 341-n_reductions sectors only
% what levels are we interested in?
n_red=[ 0 285 334 ];
%n_red=[ 291 316 331 333 ];

%Parameters values
params.rho_a = 0.9;         % persistence agg tech shocks
params.rho_mu = 0.9;         % persistence monetary shock
params.rho_k = 0.9;         % persistence idiosync tech shock
params.beta = 0.9975;      % discount factor
params.theta = 6;          % elast of subst within sectors
params.delta = 0.5;

%% Sectors
for i_k=1:length(n_red)
    kkk = n_red(i_k);
    % set the number of sectors for the economy
    n_sectors = n_sectors_full-kkk;

    % load data
    clear FPA FPA_het Omega_C Omega
    filename = ['agg_data_' num2str(n_sectors) '_C.mat']
    load(filename);

    FPA(:,:,6) = FPA_het;
    Omega_C(:,:,6) = Omega_C_het;
	Omega(:,:,6) = Omega_het;
    
    % outer loop over the parameter values, index k
for k = 1:1
    params.eta = eta_set(k);
    params.frisch = phi_set(k);
    params.sigma = sigma_set(k);
    params.phi_pi = phi_pi_set(k);
    params.phi_c = phi_c_set(k);
    params.phi_gc = phi_gc_i_set(k);
    params.rho_i = phi_rho_i_set(k);

    
    for i = 6
    i
    tic
    %%Generate Gensys matrices
        [G0, G1, C, PSI, PI, pos_c, pos_pc, pos_mc, pos_yk, pos_p, pos_ps]  = ...
        Gensys_matrix(n_sectors, FPA(:,:,i),Omega_C(:,:,i), Omega(:,:,i), params);
	toc
    	%%Call Gensys
        [G1, C, imp, fmat, fwt, ywt, gev, eu, loose] = gensys(G0, G1, C, PSI, PI);
    toc
             %check existence and uniqueness
            eu
        
            yt_imp_mu = Gensys_IRF(G1, n_sectors, n_periods, pos_mushock, imp);
            % output = [index for sectors, consumption IRF, inflation IRF, real MC IRF]
            output = [1:n_periods; yt_imp_mu(pos_c,2:end); ...
                yt_imp_mu(pos_pc,2:end) - yt_imp_mu(pos_pc,1:end-1); ...
                yt_imp_mu(pos_mc,2:end) - (yt_imp_mu(pos_pc,2:end)) ]';
            autocorr_c = autocorr(output(:,end-1),1);
            autocorr_inf = autocorr(output(:,end),1);
            autocorr_mc = autocorr([mean(yt_imp_mu(pos_mc,2:end)) - (yt_imp_mu(pos_pc,2:end))]',1);

            % output matrix: initial C, Pi, MC, cumulative C, Pi, mean(MC)...
            % AR(1) of C, Pi, MC + IRFs (C, Pi, MC)
            output_matrix = [output(:,1) repmat([output(1,2:3)...
                mean(yt_imp_mu(pos_mc,2) - (yt_imp_mu(pos_pc,2)))...
            sum(output(:,2:3)) sum(yt_imp_mu(pos_mc,2:end) - (yt_imp_mu(pos_pc,2:end)))...    
            autocorr_c(2) autocorr_inf(2) autocorr_mc(2)],n_periods,1) output(:,2:end)];

            % save output
            filename = ['Gensys_mu_output_main_case_1params1_' num2str(n_sectors_full-kkk) '.csv'];
            fid = fopen(filename, 'w');
            fprintf(fid,'Period\tConsumption_impact_irf_mu\tInflation_impact_irf_mu\tMC_mean_impact_irf_mu\tConsumption_cum_irf_mu\tInflation_cum_irf_mu\tMC_mean_cum_irf_mu\tConsumption_persistence_mu\tInflation_persistence_mu\tMC_persistence_mu\tConsumption_mu\tInflation_mu\tMC_mu \n');
            fclose(fid);
            dlmwrite(filename, (output_matrix), '-append', 'precision', '%.10f', 'delimiter', '\t');

            % output matrix: initial C, Pi, MC, cumulative C, Pi, mean(MC)...
            % AR(1) of C, Pi, MC + IRFs (C, Pi, MC)
            output_matrix = [output(:,1) repmat([output(1,2:3)...
                mean(yt_imp_mu(pos_mc,2) - (yt_imp_mu(pos_pc,2)))...
            sum(output(:,2:3)) sum(yt_imp_mu(pos_mc,2:end) - (yt_imp_mu(pos_pc,2:end)))...    
            autocorr_c(2) autocorr_inf(2) autocorr_mc(2)],n_periods,1) output(:,2:end)];

        mc_imp(:,i,1,1:n_sectors_full-kkk,n_sectors_full-kkk) = [yt_imp_mu(pos_mc,2:end)-repmat((yt_imp_mu(pos_pc,2:end)),n_sectors,1)]';

    end
    
    %%Save Case output in output matrix
    c_imp(:,i,1,n_sectors_full-kkk) = yt_imp_mu(pos_c,2:end)';
    inf_imp(:,i,1,n_sectors_full-kkk) = (yt_imp_mu(pos_pc,2:end) - yt_imp_mu(pos_pc,1:end-1))';
    
    toc
    end
end
save('alloutput_aggregation_C.mat')

%     % make figures for N-sector economies combines
i = 6
         figure('units','normalized','position',[.1 0 .4 1])
         NameArray = {'LineStyle', 'Color','Marker'}
         ValueArray = {'-','b','none'; '--','b','.'; ':','b','.'; '-.', 'b','.' ;'-','b','s'; '--','b','+';'-.','b','d'}
        % consumption
         subplot(1,1,1)
         line = plot(1:n_periods,c_imp(:,i,1,341-n_red(1))','linewidth',2)
         set(line ,NameArray, ValueArray(1,:))
    
    
         if length(n_red)>1
            hold on
    
            for ll = 2:length(n_red)
                line = plot(1:n_periods,c_imp(:,i,1,341-n_red(ll))','linewidth',2)
                set(line, NameArray, ValueArray(ll,:))
            end
         end
         
         title('\textbf{Consumption}','Interpreter','latex')
         xlabel('\textbf{Period}','Interpreter','latex')
         ylabel('\textbf{Percentage Deviation from SS}','Interpreter','latex')
           dlegend=cell(1,length(n_red));
    
        n_red_legend = n_red;
        n_red_legend(1) = 0;
        % since 349 is like 341 just relabel
        
        for hh=1:length(n_red_legend)
           dlegend(hh)={[ 'n=' num2str(341-n_red_legend(hh))]}
        end
        legend1=legend(dlegend,'Location','best')
        fig341 = gcf
        fig341.PaperUnits = 'normalized';
    %    fig341.PaperPosition = [.1 0 .8 1];
        fig341.PaperPosition = [0 0 1 .6];
        fig341.PaperPositionMode = 'manual';
        filename = ['IRFs_Nsectors_C_aggregation_C.pdf'];
        saveas(fig341,filename);
    
         % inflation
        figure()
        NameArray = {'LineStyle', 'Color','Marker'}
        ValueArray = {'-','b','none'; '--','b','.'; ':','b','.'; '-.', 'b','.' ;'-','b','s'; '--','b','+';'-.','b','d'}
        subplot(1,1,1)
        line = plot(1:n_periods,inf_imp(:,i,1,341-n_red(1))','linewidth',2)
        set(line ,NameArray, ValueArray(1,:))
    
         if length(n_red)>1
            hold on
    
            for ll = 2:length(n_red)
                line = plot(1:n_periods,inf_imp(:,i,1,341-n_red(ll))','linewidth',2)
                set(line, NameArray, ValueArray(ll,:))
            end
         end
         
        title('\textbf{Inflation}','Interpreter','latex')
        xlabel('\textbf{Period}','Interpreter','latex')
        ylabel('\textbf{Percentage Deviation from SS}','Interpreter','latex')
        dlegend=cell(1,length(n_red));
    
        n_red_legend = n_red;
        n_red_legend(1) = 0;
        
        for hh=1:length(n_red_legend)
           dlegend(hh)={[ 'n=' num2str(341-n_red_legend(hh))]}
        end
        legend1=legend(dlegend,'Location','best')
        fig341 = gcf
        fig341.PaperUnits = 'normalized';
        fig341.PaperPosition = [0 0 1 .6];
        %fig341.PaperPosition = [.1 0 .8 1];
        fig341.PaperPositionMode = 'manual';
        filename = ['IRFs_Nsectors_C_aggregation_infl.pdf'];
        saveas(fig341,filename);
    
        % marginal cost
        figure()
        subplot(1,1,1)
        NameArray = {'LineStyle', 'Color','Marker'}
        ValueArray = {'-','b','none'; '--','b','.'; ':','b','.'; '-.', 'b','.' ;'-','b','s'; '--','b','+';'-.','b','d'}
        line = plot(1:n_periods,mean(mc_imp(:,i,1,1:341-n_red(1),341-n_red(1)),4)','linewidth',2)
        set(line ,NameArray, ValueArray(1,:))
      
        if length(n_red)>1
            hold on
    
            for ll = 2:length(n_red)
                line = plot(1:n_periods,mean(mc_imp(:,i,1,1:341-n_red(ll),341-n_red(ll)),4)','linewidth',2)
                set(line, NameArray, ValueArray(ll,:))
            end
        end
        title('\textbf{Real MC}','Interpreter','latex')
        xlabel('\textbf{Period}','Interpreter','latex')
        ylabel('\textbf{Percentage Deviation from SS}','Interpreter','latex')
        dlegend=cell(1,length(n_red));
    
        n_red_legend = n_red;
        n_red_legend(1) = 0;
        % since 349 is like 341 just relabel
        
        for hh=1:length(n_red_legend)
           dlegend(hh)={[ 'n=' num2str(341-n_red_legend(hh))]}
        end
        legend1=legend(dlegend,'Location','best')
        fig341 = gcf
        fig341.PaperUnits = 'normalized';
        fig341.PaperPosition = [0 0 1 .6];
    %    fig341.PaperPosition = [.1 0 .8 1];
        fig341.PaperPositionMode = 'manual';
        filename = ['IRFs_Nsectors_C_aggregation_mc.pdf'];
        saveas(fig341,filename);
    
        %compute the persistence of each aggregation level
        for ll=1:length(n_red)
            xxxx=autocorr(c_imp(:,i,1,341-n_red(ll)),1);
            rho_cons(ll)=xxxx(2);
            xxxx=autocorr(inf_imp(:,i,1,341-n_red(ll)),1);
            rho_pi(ll)=xxxx(2);
        end
        disp('Autocorrelation of consumption and inflation:')
        disp([rho_cons' rho_pi'])
